from flask import Flask, render_template, request
from Crypto.Cipher import AES, DES, DES3
from Crypto.Util.Padding import pad, unpad
import base64
import hashlib
from Crypto.Random import get_random_bytes
from flask import send_from_directory

app = Flask(__name__)

# AES加密
def pad_key(key, required_length):
    """自动填充密钥到指定长度"""
    if len(key) >= required_length:
        return key[:required_length]
    # 重复密钥直到达到所需长度
    return (key * (required_length // len(key) + 1))[:required_length]

def encrypt_aes(plaintext, key):
    # 检查并调整密钥长度
    if len(key) < 16:
        key = pad_key(key, 16)  # AES-128
    elif len(key) < 24:
        key = pad_key(key, 24)  # AES-192
    elif len(key) < 32:
        key = pad_key(key, 32)  # AES-256
    else:
        key = key[:32]  # 截断过长的密钥
    
    cipher = AES.new(key.encode('utf-8'), AES.MODE_ECB)
    ct_bytes = cipher.encrypt(pad(plaintext.encode('utf-8'), AES.block_size))
    return base64.b64encode(ct_bytes).decode('utf-8')

# AES解密
def decrypt_aes(ciphertext, key):
    # 使用与加密相同的密钥处理逻辑
    if len(key) < 16:
        key = pad_key(key, 16)
    elif len(key) < 24:
        key = pad_key(key, 24)
    elif len(key) < 32:
        key = pad_key(key, 32)
    else:
        key = key[:32]
    
    cipher = AES.new(key.encode('utf-8'), AES.MODE_ECB)
    ct_bytes = base64.b64decode(ciphertext)
    pt = unpad(cipher.decrypt(ct_bytes), AES.block_size)
    return pt.decode('utf-8')

# DES加密
def encrypt_des(plaintext, key):
    cipher = DES.new(key.encode('utf-8'), DES.MODE_ECB)
    ct_bytes = cipher.encrypt(pad(plaintext.encode('utf-8'), DES.block_size))
    return base64.b64encode(ct_bytes).decode('utf-8')

# DES解密
def decrypt_des(ciphertext, key):
    cipher = DES.new(key.encode('utf-8'), DES.MODE_ECB)
    ct_bytes = base64.b64decode(ciphertext)
    pt = unpad(cipher.decrypt(ct_bytes), DES.block_size)
    return pt.decode('utf-8')

# Triple DES加密
def encrypt_3des(plaintext, key):
    cipher = DES3.new(key.encode('utf-8'), DES3.MODE_ECB)
    ct_bytes = cipher.encrypt(pad(plaintext.encode('utf-8'), DES3.block_size))
    return base64.b64encode(ct_bytes).decode('utf-8')

# Triple DES解密
def decrypt_3des(ciphertext, key):
    cipher = DES3.new(key.encode('utf-8'), DES3.MODE_ECB)
    ct_bytes = base64.b64decode(ciphertext)
    pt = unpad(cipher.decrypt(ct_bytes), DES3.block_size)
    return pt.decode('utf-8')

# RC4加密/解密
def rc4_crypt(text, key):
    S = list(range(256))
    j = 0
    for i in range(256):
        j = (j + S[i] + ord(key[i % len(key)])) % 256
        S[i], S[j] = S[j], S[i]
    i = j = 0
    result = []
    for char in text:
        i = (i + 1) % 256
        j = (j + S[i]) % 256
        S[i], S[j] = S[j], S[i]
        result.append(chr(ord(char) ^ S[(S[i] + S[j]) % 256]))
    return ''.join(result)

# Rabbit加密/解密
def rabbit_crypt(text, key):
    # 简化的Rabbit实现
    key = key.encode('utf-8')
    iv = get_random_bytes(8)
    cipher = AES.new(key, AES.MODE_CFB, iv)
    return base64.b64encode(cipher.encrypt(text.encode('utf-8'))).decode('utf-8')

# MD5哈希
def md5_hash(text):
    return hashlib.md5(text.encode('utf-8')).hexdigest()

# Base64编码
def base64_encode(text):
    return base64.b64encode(text.encode('utf-8')).decode('utf-8')

# Base64解码
def base64_decode(text):
    return base64.b64decode(text).decode('utf-8')

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/encrypt', methods=['GET', 'POST'])
def encrypt():
    if request.method == 'POST':
        plaintext = request.form['plaintext']
        key = request.form['key']
        algorithm = request.form['algorithm']
        
        if algorithm == 'aes':
            ciphertext = encrypt_aes(plaintext, key)
        elif algorithm == 'des':
            ciphertext = encrypt_des(plaintext, key)
        elif algorithm == '3des':
            ciphertext = encrypt_3des(plaintext, key)
        elif algorithm == 'rc4':
            ciphertext = rc4_crypt(plaintext, key)
        elif algorithm == 'rabbit':
            ciphertext = rabbit_crypt(plaintext, key)
        elif algorithm == 'md5':
            ciphertext = md5_hash(plaintext)
        elif algorithm == 'base64':
            ciphertext = base64_encode(plaintext)
        else:
            ciphertext = "不支持的算法"
            
        return render_template('encrypt.html', ciphertext=ciphertext)
    return render_template('encrypt.html')

@app.route('/decrypt', methods=['GET', 'POST'])
def decrypt():
    if request.method == 'POST':
        ciphertext = request.form['ciphertext']
        key = request.form['key']
        algorithm = request.form['algorithm']
        
        try:
            if algorithm == 'aes':
                plaintext = decrypt_aes(ciphertext, key)
            elif algorithm == 'des':
                plaintext = decrypt_des(ciphertext, key)
            elif algorithm == '3des':
                plaintext = decrypt_3des(ciphertext, key)
            elif algorithm == 'rc4':
                plaintext = rc4_crypt(ciphertext, key)
            elif algorithm == 'rabbit':
                plaintext = rabbit_crypt(ciphertext, key)
            elif algorithm == 'md5':
                plaintext = "MD5验证: " + ("匹配" if md5_hash(ciphertext) == key else "不匹配")
            elif algorithm == 'base64':
                plaintext = base64_decode(ciphertext)
            else:
                plaintext = "不支持的算法或无法解密"
        except:
            plaintext = "解密失败，请检查密钥和密文是否正确"
            
        return render_template('decrypt.html',
                            plaintext=plaintext,
                            ord=ord)
    return render_template('decrypt.html', ord=ord)


# 添加静态文件路由
@app.route('/static/<path:filename>')
def static_files(filename):
    return send_from_directory('static', filename)

if __name__ == '__main__':
    app.run(debug=True)